import pandas as pd
import numpy as np
from prophet import Prophet
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta
import warnings
import pickle
import os
import sys
import argparse
from collections import defaultdict
from Noise import *


# --- 路径设置 ---
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_DIR = os.path.dirname(SCRIPT_DIR)
DATA_DIR = os.path.join(BASE_DIR, 'data')
MODELS_DIR = os.path.join(BASE_DIR, 'models')
REPORTS_DIR = os.path.join(BASE_DIR, 'reports')
os.makedirs(MODELS_DIR, exist_ok=True)
os.makedirs(REPORTS_DIR, exist_ok=True)

# --- 日志记录 ---
class Logger(object):
    def __init__(self, filename=os.path.join(REPORTS_DIR, "enhanced_training_log.log")):
        self.terminal = sys.stdout
        self.log = open(filename, "w", encoding="utf-8")
    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)
    def flush(self):
        pass
sys.stdout = Logger()

# --- 全局设置 ---
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
warnings.filterwarnings('ignore')

class EnhancedInventoryProphetTrainer:
    def __init__(self, data_filename='processed_new_data.csv', seasonality_mode='multiplicative', 
                 rolling_window=7, safety_stock_factor=1.2, min_product_records=10):
        self.data_path = os.path.join(DATA_DIR, data_filename)
        self.seasonality_mode = seasonality_mode
        self.rolling_window = rolling_window
        self.safety_stock_factor = safety_stock_factor
        self.min_product_records = min_product_records
        self.data = None
        self.processed_data = None
        self.product_models = {}  # 存储每个商品的模型
        self.aggregate_models = {}  # 存储聚合模型
        self.results = {}
        self.product_info = {}  # 存储商品信息
        
        print("--- 增强版训练器初始化 ---")
        print(f"  - 数据文件: {data_filename}")
        print(f"  - 季节性模式: {self.seasonality_mode}")
        print(f"  - 移动平均窗口: {self.rolling_window}")
        print(f"  - 安全库存系数: {self.safety_stock_factor}")
        print(f"  - 最小商品记录数: {self.min_product_records}")
        
    def load_data(self):
        print("正在加载数据...")
        try:
            # 尝试加载预处理后的数据
            if self.data_path.endswith('.csv'):
                self.data = pd.read_csv(self.data_path, dtype={"商品ID": str})
            else:
                self.data = pd.read_excel(self.data_path, header=0, dtype={"商品ID": str})
            print(f"数据加载成功！数据形状: {self.data.shape}")
            self._validate_data_structure()
            return True
        except Exception as e:
            print(f"数据加载失败: {e}")
            return False
    
    def _validate_data_structure(self):
        """验证数据结构并输出详细信息"""
        print("\n=== 数据结构分析 ===")
        print(f"原始数据形状: {self.data.shape}")
        
        # 检查关键列是否存在
        required_columns = ['统计日期', '商品ID', '支付件数', '支付金额']
        missing_columns = [col for col in required_columns if col not in self.data.columns]
        if missing_columns:
            print(f"警告: 缺少必需列: {missing_columns}")
            print(f"可用列: {list(self.data.columns)}")
            return False
        
        # 先转换日期字段，避免后续比较错误
        try:
            if self.data['统计日期'].dtype == 'object':
                self.data['统计日期'] = pd.to_datetime(self.data['统计日期'], errors='coerce')
            # 移除无效日期
            invalid_dates = self.data['统计日期'].isna().sum()
            if invalid_dates > 0:
                print(f"警告: 发现 {invalid_dates} 条无效日期记录，将被过滤")
                self.data = self.data.dropna(subset=['统计日期'])
        except Exception as e:
            print(f"日期字段转换失败: {e}")
            return False
        
        # 分析数据分布
        print(f"商品数量: {self.data['商品ID'].nunique()}")
        print(f"日期范围: {self.data['统计日期'].min()} 到 {self.data['统计日期'].max()}")
        print(f"总记录数: {len(self.data)}")
        
        # 检查每个商品的记录数
        product_counts = self.data['商品ID'].value_counts()
        print(f"商品记录数统计:")
        print(f"  最少记录: {product_counts.min()}")
        print(f"  最多记录: {product_counts.max()}")
        print(f"  平均记录: {product_counts.mean():.2f}")
        print(f"  中位数记录: {product_counts.median()}")
        
        # 统计有足够记录的商品数量
        sufficient_products = (product_counts >= self.min_product_records).sum()
        print(f"有足够记录的商品数量 (≥{self.min_product_records}条): {sufficient_products}")
        
        # 检查新数据字段
        new_fields = ['商品访客数', '商品浏览量', '商品平均停留时长', '商品详情页跳出率', 
                     '商品加购人数', '商品加购件数', '商品收藏人数', '下单转换率', 
                     '支付转换率', '访客平均价值']
        available_new_fields = [field for field in new_fields if field in self.data.columns]
        if available_new_fields:
            print(f"新数据字段: {available_new_fields}")
        
        return True
    
    def calculate_inventory_demand(self):
        print("\n正在计算进货需求指标...")
        if self.data is None: return False
        
        df = self.data.copy()
        
        print(f"数据清洗后形状: {df.shape}")
        
        # 1. 保留商品粒度的数据
        print("\n=== 保留商品粒度处理 ===")
        daily_product_data = df.groupby(['统计日期', '商品ID']).agg({
            '支付件数': 'sum',
            '支付金额': 'sum',
            '商品访客数': 'sum' if '商品访客数' in df.columns else lambda x: 0,
            '商品浏览量': 'sum' if '商品浏览量' in df.columns else lambda x: 0,
            '商品加购件数': 'sum' if '商品加购件数' in df.columns else lambda x: 0,
            '商品收藏人数': 'sum' if '商品收藏人数' in df.columns else lambda x: 0
        }).reset_index()
        
        # 添加商品信息（如果存在）
        if '商品名称' in df.columns:
            product_info = df.groupby('商品ID').agg({
                '商品名称': 'first',
                '货号': 'first' if '货号' in df.columns else lambda x: 'N/A'
            }).reset_index()
            daily_product_data = daily_product_data.merge(product_info, on='商品ID', how='left')
        
        print(f"商品粒度数据形状: {daily_product_data.shape}")
        
        # 2. 创建聚合数据（用于整体趋势预测）
        print("\n=== 创建聚合数据 ===")
        aggregate_data = daily_product_data.groupby('统计日期').agg({
            '支付件数': 'sum',
            '支付金额': 'sum',
            '商品访客数': 'sum' if '商品访客数' in daily_product_data.columns else lambda x: 0,
            '商品浏览量': 'sum' if '商品浏览量' in daily_product_data.columns else lambda x: 0,
            '商品加购件数': 'sum' if '商品加购件数' in daily_product_data.columns else lambda x: 0,
            '商品收藏人数': 'sum' if '商品收藏人数' in daily_product_data.columns else lambda x: 0
        }).reset_index()
        
        # 过滤无效的零销售期
        active_sales_df = aggregate_data[aggregate_data['支付件数'] > 0]
        if not active_sales_df.empty:
            first_sale_date = active_sales_df['统计日期'].min()
            last_sale_date = active_sales_df['统计日期'].max()
            
            aggregate_data = aggregate_data[
                (aggregate_data['统计日期'] >= first_sale_date) & 
                (aggregate_data['统计日期'] <= last_sale_date)
            ]
            print(f"有效销售周期: {first_sale_date.strftime('%Y-%m-%d')} 到 {last_sale_date.strftime('%Y-%m-%d')}")
        
        # 3. 计算需求预测指标
        aggregate_data['需求预测'] = aggregate_data['支付件数'].rolling(
            window=self.rolling_window, min_periods=1
        ).mean() * self.safety_stock_factor
        
        # 4. 为每个商品计算需求预测
        print("\n=== 计算商品级别需求预测 ===")
        for product_id in daily_product_data['商品ID'].unique():
            product_data = daily_product_data[daily_product_data['商品ID'] == product_id].copy()
            if len(product_data) >= self.min_product_records:
                product_data['需求预测'] = product_data['支付件数'].rolling(
                    window=self.rolling_window, min_periods=1
                ).mean() * self.safety_stock_factor
                daily_product_data.loc[daily_product_data['商品ID'] == product_id, '需求预测'] = product_data['需求预测']
        
        # 5. 重命名列以适配Prophet
        aggregate_data.rename(columns={'统计日期': 'ds', '支付件数': 'sales'}, inplace=True)
        daily_product_data.rename(columns={'统计日期': 'ds', '支付件数': 'sales'}, inplace=True)
        
        # 存储处理后的数据
        self.processed_data = {
            'aggregate': aggregate_data,
            'product_level': daily_product_data
        }
        
        print(f"聚合数据形状: {aggregate_data.shape}")
        print(f"商品级别数据形状: {daily_product_data.shape}")
        print("数据预处理完成！")
        return True
    
    def train_models(self):
        """训练聚合模型和商品级别模型"""
        if not self.processed_data: return False
        
        print("\n=== 开始训练模型 ===")
        
        # 1. 训练聚合模型
        print("\n--- 训练聚合模型 ---")
        self._train_aggregate_models()
        
        # 2. 训练商品级别模型
        print("\n--- 训练商品级别模型 ---")
        self._train_product_models()
        
        return len(self.aggregate_models) > 0 or len(self.product_models) > 0
    
    def _train_aggregate_models(self):
        """训练聚合数据模型"""
        aggregate_data = self.processed_data['aggregate']
        target_columns = ['sales', '需求预测']
        
        for target in target_columns:
            print(f"\n训练聚合模型: {target}")
            model_data = aggregate_data[['ds', target]].copy().rename(columns={target: 'y'})
            model_data.dropna(inplace=True)
            
            if len(model_data) < 10:
                print(f"  数据点不足 ({len(model_data)} < 10)，跳过")
                continue
                
            try:
                model = Prophet(
                    yearly_seasonality=False,  # 数据量少于一年，关闭年度季节性
                    weekly_seasonality=False,   # 数据包含7个完整周，保留周度季节性
                    daily_seasonality=False, # 日度数据，关闭日度季节性
                    seasonality_mode=self.seasonality_mode
                )
                # 添加30天月周期季节性
                model.add_seasonality(
                    name='monthly',  # 自定义名称
                    period=30,  # 周期长度（天）
                    fourier_order=5  # 傅里叶级数阶数
                )

                # 如果有特定的节假日信息，可以在这里添加
                model.add_country_holidays(country_name='CN')

                # 生成女性衣物专属购物节
                fashion_holidays = generate_women_fashion_holidays()

                # 添加女性专属购物节
                model.holidays = pd.concat([model.holidays, fashion_holidays])
                # 添加天气相关因子（服装销售关键）
                # model.add_regressor('temperature')
                # model.add_regressor('precipitation')

                model.fit(model_data)
                self.aggregate_models[target] = model
                self._evaluate_model(model, model_data, f"aggregate_{target}")
                print(f"  聚合模型 {target} 训练完成！")
            except Exception as e:
                print(f"  训练聚合模型 {target} 时出错: {e}")
    
    def _train_product_models(self):
        """训练商品级别模型"""
        product_data = self.processed_data['product_level']
        product_ids = product_data['商品ID'].unique()
        
        successful_products = 0
        total_products = len(product_ids)
        
        for i, product_id in enumerate(product_ids, 1):
            product_info = product_data[product_data['商品ID'] == product_id].iloc[0]
            product_name = product_info.get('商品名称', f'商品_{product_id}')
            
            print(f"\n训练商品模型 [{i}/{total_products}]: {product_name} (ID: {product_id})")
            
            # 获取该商品的数据
            product_model_data = product_data[product_data['商品ID'] == product_id].copy()
            
            if len(product_model_data) < self.min_product_records:
                print(f"  记录数不足 ({len(product_model_data)} < {self.min_product_records})，跳过")
                continue
            
            # 训练销售预测模型
            try:
                sales_data = product_model_data[['ds', 'sales']].copy().rename(columns={'sales': 'y'})
                sales_data.dropna(inplace=True)
                
                if len(sales_data) >= 30:
                    model = Prophet(
                        yearly_seasonality=False,
                        weekly_seasonality=False,
                        daily_seasonality=False,
                        # seasonality_mode='multiplicative',  # 服装销售受季节比例影响
                        changepoint_prior_scale=0.03,  # 减少趋势突变敏感度
                        holidays_prior_scale=12.0, # 增强节日影响
                        seasonality_mode=self.seasonality_mode
                    )
                    # 添加30天月周期季节性
                    model.add_seasonality(
                        name='monthly',  # 自定义名称
                        period=30,  # 周期长度（天）
                        fourier_order=5  # 傅里叶级数阶数
                    )

                    # 如果有特定的节假日信息，可以在这里添加
                    model.add_country_holidays(country_name='CN')

                    # 生成女性衣物专属购物节
                    fashion_holidays = generate_women_fashion_holidays()

                    # 添加女性专属购物节
                    model.holidays = pd.concat([model.holidays, fashion_holidays])
                    # 添加天气相关因子（服装销售关键）
                    # model.add_regressor('temperature')
                    # model.add_regressor('precipitation')

                    model.fit(sales_data)

                    # 存储模型和商品信息
                    self.product_models[product_id] = {
                        'sales_model': model,
                        'product_name': product_name,
                        'product_info': product_info.to_dict(),
                        'data_points': len(sales_data)
                    }
                    
                    # 评估模型
                    self._evaluate_model(model, sales_data, f"product_{product_id}_sales")
                    successful_products += 1
                    print(f"  商品 {product_name} 销售模型训练完成！")
                else:
                    print(f"  有效数据点不足 ({len(sales_data)} < 10)，跳过")
                    
            except Exception as e:
                print(f"  训练商品 {product_name} 模型时出错: {e}")
        
        print(f"\n商品级别模型训练完成！成功训练 {successful_products}/{total_products} 个商品模型")
    
    def _evaluate_model(self, model, data, target_name):
        """评估模型性能"""
        try:
            forecast = model.predict(data)
            actual = data['y'].values
            predicted = forecast['yhat'].values
            
            # 计算评估指标
            r2 = r2_score(actual, predicted)
            mae = mean_absolute_error(actual, predicted)
            rmse = np.sqrt(mean_squared_error(actual, predicted))
            mape = np.mean(np.abs((actual - predicted) / (actual + 1e-8))) * 100
            
            self.results[target_name] = {
                'r2': r2, 
                'mae': mae, 
                'rmse': rmse, 
                'mape': mape,
                'forecast': forecast
            }
            
            print(f"  R²: {r2:.4f}, MAE: {mae:.2f}, RMSE: {rmse:.2f}, MAPE: {mape:.2f}%")
            
        except Exception as e:
            print(f"  评估模型时出错: {e}")
    
    def save_models(self):
        """保存所有模型"""
        print("\n=== 保存模型 ===")
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 保存聚合模型
        if self.aggregate_models:
            aggregate_filename = f"aggregate_prophet_models_{timestamp}.pkl"
            aggregate_filepath = os.path.join(MODELS_DIR, aggregate_filename)
            
            aggregate_payload = {
                'models': self.aggregate_models,
                'results': {k: v for k, v in self.results.items() if k.startswith('aggregate_')},
                'metadata': {
                    'timestamp': timestamp,
                    'model_type': 'aggregate',
                    'seasonality_mode': self.seasonality_mode,
                    'rolling_window': self.rolling_window,
                    'safety_stock_factor': self.safety_stock_factor
                }
            }
            
            with open(aggregate_filepath, 'wb') as f:
                pickle.dump(aggregate_payload, f)
            print(f"聚合模型已保存到: {aggregate_filepath}")
        
        # 保存商品级别模型
        if self.product_models:
            product_filename = f"product_prophet_models_{timestamp}.pkl"
            product_filepath = os.path.join(MODELS_DIR, product_filename)
            
            product_payload = {
                'models': self.product_models,
                'results': {k: v for k, v in self.results.items() if k.startswith('product_')},
                'metadata': {
                    'timestamp': timestamp,
                    'model_type': 'product_level',
                    'seasonality_mode': self.seasonality_mode,
                    'rolling_window': self.rolling_window,
                    'safety_stock_factor': self.safety_stock_factor,
                    'min_product_records': self.min_product_records
                }
            }
            
            with open(product_filepath, 'wb') as f:
                pickle.dump(product_payload, f)
            print(f"商品级别模型已保存到: {product_filepath}")
        
        # 保存最佳模型（用于Web应用）
        if self.aggregate_models:
            best_target = max(
                [k for k in self.results.keys() if k.startswith('aggregate_')],
                key=lambda x: self.results[x]['r2']
            )
            best_score = self.results[best_target]['r2']
            best_model = self.aggregate_models[best_target.replace('aggregate_', '')]
            
            best_filename = f"best_prophet_model_accuracy_{best_score:.4f}_{timestamp}.pkl"
            best_filepath = os.path.join(MODELS_DIR, best_filename)
            
            best_payload = {
                'model': best_model,
                'target': best_target.replace('aggregate_', ''),
                'metadata': {
                    'accuracy': best_score,
                    'timestamp': timestamp,
                    'model_type': 'best_aggregate'
                }
            }
            
            with open(best_filepath, 'wb') as f:
                pickle.dump(best_payload, f)
            print(f"最佳模型已保存到: {best_filepath}")
    
    def generate_enhanced_report(self):
        """生成训练报告"""
        print("\n=== 生成训练报告 ===")
        
        report_content = []
        report_content.append("=" * 60)
        report_content.append("库存需求预测Prophet模型训练报告")
        report_content.append("=" * 60)
        report_content.append(f"训练时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        report_content.append(f"数据文件: {os.path.basename(self.data_path)}")
        report_content.append(f"原始数据形状: {self.data.shape}")
        
        if self.processed_data:
            report_content.append(f"聚合数据形状: {self.processed_data['aggregate'].shape}")
            report_content.append(f"商品级别数据形状: {self.processed_data['product_level'].shape}")
        
        report_content.append("")
        report_content.append("聚合模型性能:")
        report_content.append("-" * 40)
        
        aggregate_results = {k: v for k, v in self.results.items() if k.startswith('aggregate_')}
        for target, metrics in aggregate_results.items():
            target_name = target.replace('aggregate_', '')
            report_content.append(f"{target_name}:")
            report_content.append(f"  R²: {metrics['r2']:.4f}")
            report_content.append(f"  MAE: {metrics['mae']:.2f}")
            report_content.append(f"  RMSE: {metrics['rmse']:.2f}")
            report_content.append(f"  MAPE: {metrics['mape']:.2f}%")
            report_content.append("")
        
        report_content.append("商品级别模型统计:")
        report_content.append("-" * 40)
        report_content.append(f"成功训练的商品模型数量: {len(self.product_models)}")
        
        if self.product_models:
            # 统计商品模型性能
            product_r2_scores = []
            for product_id, model_info in self.product_models.items():
                result_key = f"product_{product_id}_sales"
                if result_key in self.results:
                    product_r2_scores.append(self.results[result_key]['r2'])
            
            if product_r2_scores:
                report_content.append(f"商品模型R²分数统计:")
                report_content.append(f"  平均R²: {np.mean(product_r2_scores):.4f}")
                report_content.append(f"  最高R²: {np.max(product_r2_scores):.4f}")
                report_content.append(f"  最低R²: {np.min(product_r2_scores):.4f}")
                report_content.append(f"  中位数R²: {np.median(product_r2_scores):.4f}")
        
        report_content.append("")
        report_content.append("改进说明:")
        report_content.append("-" * 40)
        report_content.append("1. 保留了商品粒度信息，支持商品级别预测")
        report_content.append("2. 同时提供聚合模型和商品级别模型")
        report_content.append("3. 增加了数据验证和详细日志")
        report_content.append("4. 支持最小记录数过滤，确保模型质量")
        report_content.append("5. 提供更全面的模型性能评估")
        
        # 保存报告
        report_filename = f"enhanced_inventory_prophet_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
        report_filepath = os.path.join(REPORTS_DIR, report_filename)
        
        with open(report_filepath, 'w', encoding='utf-8') as f:
            f.write('\n'.join(report_content))
        
        print(f"训练报告已保存到: {report_filepath}")
        
        # 打印报告摘要
        print("\n" + "=" * 60)
        print("训练报告摘要:")
        print("=" * 60)
        for line in report_content[:20]:  # 只打印前20行
            print(line)

def main(args):
    print("--- 开始Prophet模型训练流程 ---")
    trainer = EnhancedInventoryProphetTrainer(
        data_filename=args.data_file,
        seasonality_mode=args.seasonality_mode,
        rolling_window=args.rolling_window,
        safety_stock_factor=args.safety_stock_factor,
        min_product_records=args.min_product_records
    )
    
    if trainer.load_data():
        if trainer.calculate_inventory_demand():
            if trainer.train_models():
                trainer.save_models()
                trainer.generate_enhanced_report()
    
    print("\n--- 增强版训练流程结束 ---")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="库存预测Prophet模型训练器")
    
    parser.add_argument('--data_file', 
                        type=str, 
                        default='processed_new_data.csv', 
                        help='位于 data/ 目录下的源数据文件名')
                        
    parser.add_argument('--seasonality_mode', 
                        type=str, 
                        default='multiplicative', 
                        choices=['additive', 'multiplicative'], 
                        help='Prophet模型的季节性模式')
                        
    parser.add_argument('--rolling_window', 
                        type=int, 
                        default=7, 
                        help='移动平均窗口大小（天数）')
                        
    parser.add_argument('--safety_stock_factor', 
                        type=float, 
                        default=1.2, 
                        help='安全库存系数')
    
    parser.add_argument('--min_product_records', 
                        type=int, 
                        default=10, 
                        help='每个商品的最小记录数要求')

    args = parser.parse_args()
    main(args) 